package com.yhg.backend.config;

import jakarta.servlet.http.HttpSession;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * WebSocket握手拦截器，用于在WebSocket握手前确保会话已建立
 */
@Component
public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 如果是ServletServerHttpRequest，则获取HttpSession
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            HttpSession session = servletRequest.getServletRequest().getSession(false);
            
            if (session != null) {
                // 将HttpSession中的用户信息放入WebSocket的attributes中
                String currentUser = (String) session.getAttribute("currentUser");
                if (currentUser != null) {
                    System.out.println("WebSocket握手拦截器: 从HttpSession获取到用户 " + currentUser);
                    attributes.put("currentUser", currentUser);
                    return true;
                }
            }
            
            // 如果没有从HttpSession获取到用户信息，尝试从URL参数获取
            String username = servletRequest.getServletRequest().getParameter("username");
            if (username != null && !username.isEmpty()) {
                System.out.println("WebSocket握手拦截器: 从URL参数获取到用户 " + username);
                attributes.put("currentUser", username);
                return true;
            }
            
            System.out.println("WebSocket握手拦截器: 无法获取用户信息，握手失败");
            return false;
        }
        
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        // 握手后的处理，可以为空
        if (exception != null) {
            System.err.println("WebSocket握手后发生异常: " + exception.getMessage());
        }
    }
}